from langchain_core.documents import Document
from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain.text_splitter import TokenTextSplitter
from langchain_community.document_loaders import PyPDFium2Loader
from pathlib import Path
import warnings
from typing import List, Dict
import re
import logging
import unicodedata
from zhipuai import ZhipuAI

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# 配置参数
ZHIPU_API_KEY = "2f39319bdd864fc4a41bf6b8eed6efbc.uIsAkRrMwejVTIyc"
INITIAL_DOC_PATH = '01_“未来校园”智能应用专项赛.pdf'
MAX_TEXT_LENGTH = 4000  # 根据智谱API调整

# 初始化智谱客户端
zhipu_client = ZhipuAI(api_key=ZHIPU_API_KEY)


# 增强型文本预处理模块
class TextPreprocessor:
    @staticmethod
    def clean_text(text: str) -> str:
        """深度文本清洗"""
        # 1. 去除控制字符
        text = re.sub(r'[\x00-\x1F\x7F-\x9F]', '', text)
        # 2. 标准化Unicode字符
        text = unicodedata.normalize('NFKC', text)
        # 3. 替换全角字符
        fullwidth_chars = '　‘’“”【】·―—（）％＃＠＆'
        halfwidth_chars = ' \'\'\"\"[]-—()%#@&'
        trans_table = str.maketrans(fullwidth_chars, halfwidth_chars)
        text = text.translate(trans_table)
        # 4. 清理空白字符
        text = re.sub(r'\s+', ' ', text).strip()
        # 5. 数学符号转换
        math_symbols = {'≤': '<=', '≥': '>=', '≠': '!=', '±': '+/-'}
        text = re.sub(r'([\u2200-\u22FF])', lambda m: math_symbols.get(m.group(1), m.group(1)), text)
        # 6. 截断至模型限制
        return text[:MAX_TEXT_LENGTH]

    @classmethod
    def process_document(cls, doc: Document) -> Document:
        """处理单个文档"""
        cleaned_content = cls.clean_text(doc.page_content)
        if len(cleaned_content) < 50:
            raise ValueError("有效文本过短（可能为扫描件或空白页）")
        return Document(page_content=cleaned_content, metadata=doc.metadata)


# 初始化文本分割器
text_splitter = TokenTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    encoding_name="cl100k_base"
)

# 初始化嵌入模型
embeddings = ZhipuAIEmbeddings(api_key=ZHIPU_API_KEY)


# 初始化知识库
def initialize_vectorstore() -> Qdrant:
    """初始化向量存储"""
    try:
        loader = PyPDFium2Loader(INITIAL_DOC_PATH, render_pages=True)
        raw_docs = loader.load()

        processed_docs = []
        for doc in raw_docs:
            try:
                processed_docs.append(TextPreprocessor.process_document(doc))
            except Exception as e:
                logger.warning(f"初始文档处理跳过页 {doc.metadata.get('page')}: {str(e)}")

        chunks = text_splitter.split_documents(processed_docs)
        logger.info(f"初始文档分割为 {len(chunks)} 个有效块")

        return Qdrant.from_documents(
            documents=chunks,
            embedding=embeddings,
            location=":memory:",
            collection_name="smart_campus_docs"
        )
    except Exception as e:
        logger.error(f"知识库初始化失败: {str(e)}")
        raise


vectorstore = initialize_vectorstore()


# 增强版文档处理
def process_file(file_path: Path) -> Dict:
    result = {
        'file': str(file_path),
        'status': 'success',
        'chunks': 0,
        'errors': []
    }

    try:
        # 文件验证
        if not file_path.exists():
            raise FileNotFoundError(f"文件不存在: {file_path}")
        if file_path.suffix.lower() != '.pdf':
            raise ValueError("仅支持PDF格式")

        # 加载PDF内容（增强渲染模式）
        loader = PyPDFium2Loader(str(file_path), render_pages=True)
        raw_docs = loader.load_and_split()

        # 多阶段处理
        valid_chunks = []
        for doc in raw_docs:
            try:
                processed_doc = TextPreprocessor.process_document(doc)
                chunks = text_splitter.split_documents([processed_doc])
                valid_chunks.extend(chunks)
            except Exception as page_error:
                result['errors'].append({
                    'page': doc.metadata.get('page', 'unknown'),
                    'error': str(page_error),
                    'sample': doc.page_content[:100] + '...' if doc.page_content else '[空内容]'
                })

        if not valid_chunks:
            raise ValueError("无有效文本内容")

        # 分批提交（带API验证）
        batch_size = 5
        success_count = 0
        for i in range(0, len(valid_chunks), batch_size):
            batch = valid_chunks[i:i + batch_size]
            try:
                # API参数验证
                for chunk in batch:
                    if len(chunk.page_content.encode('utf-8')) > 4096:
                        raise ValueError(f"分块字节长度超标 ({len(chunk.page_content.encode('utf-8'))} > 4096)")

                vectorstore.add_documents(batch)
                success_count += len(batch)
            except Exception as batch_error:
                result['errors'].append({
                    'batch': f"{i}-{i + batch_size}",
                    'error': str(batch_error)[:200]
                })
                logger.error(f"分块提交失败: {str(batch_error)}")

        result['chunks'] = success_count
        logger.info(f"成功加载 {file_path.name} ({success_count}/{len(valid_chunks)} 分块)")

    except Exception as e:
        result['status'] = 'failed'
        result['errors'].append({'global': str(e)})
        logger.error(f"文件处理失败: {str(e)}")

    return result


# 批量处理函数（日志集成版）
def batch_add_documents(paths: List[str]) -> Dict:
    """批量处理文档"""
    total_results = {
        'total_files': 0,
        'success_files': 0,
        'failed_files': 0,
        'total_chunks': 0,
        'errors': []
    }

    processed_files = set()

    for input_path in paths:
        current_path = Path(input_path.strip()).resolve()

        # 处理文件夹
        if current_path.is_dir():
            logger.info(f"扫描文件夹: {current_path}")
            for file_path in current_path.glob('**/*.pdf'):
                if file_path in processed_files:
                    continue

                total_results['total_files'] += 1
                res = process_file(file_path)
                processed_files.add(file_path)

                if res['status'] == 'success':
                    total_results['success_files'] += 1
                    total_results['total_chunks'] += res['chunks']
                else:
                    total_results['failed_files'] += 1
                    total_results['errors'].append({
                        'file': res['file'],
                        'errors': res['errors']
                    })

        # 处理单个文件
        elif current_path.is_file():
            if current_path in processed_files:
                continue

            total_results['total_files'] += 1
            res = process_file(current_path)
            processed_files.add(current_path)

            if res['status'] == 'success':
                total_results['success_files'] += 1
                total_results['total_chunks'] += res['chunks']
            else:
                total_results['failed_files'] += 1
                total_results['errors'].append({
                    'file': res['file'],
                    'errors': res['errors']
                })

        else:
            error_msg = f"路径不存在: {current_path}"
            logger.error(error_msg)
            total_results['errors'].append({
                'file': str(current_path),
                'error': error_msg
            })

    return total_results


# 构建问答链
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
qa_chain = (
        RunnableParallel({
            "context": retriever,
            "question": RunnablePassthrough()
        })
        | ChatPromptTemplate.from_template("""
        基于以下上下文提供准确回答：
        {context}

        问题：{question}

        若上下文无相关信息，请回答“根据现有资料无法回答”
        """)
        | ChatZhipuAI(
    api_key=ZHIPU_API_KEY,
    temperature=0.3,
    model="glm-4",
    max_tokens=1000
)
        | StrOutputParser()
)


# 交互界面（增强版）
def interactive_interface():
    """交互式界面"""
    print("\n" + "=" * 40)
    print("🏫 智能校园知识助手 GLM-4")
    print("=" * 40)
    print("功能菜单:")
    print("1. 提问\n2. 批量添加文档\n3. 查看知识库状态\n4. 退出")

    while True:
        try:
            choice = input("\n请输入选项：").strip()

            # 提问模式
            if choice == '1':
                query = input("\n请输入问题：").strip()
                if not query:
                    print("⚠️ 问题不能为空")
                    continue

                print("\n🔍 正在生成回答...")
                try:
                    response = qa_chain.invoke(query)
                    print("\n" + "-" * 40)
                    print("🤖 回答：")
                    print(response)
                    print("-" * 40)
                except Exception as e:
                    logger.error(f"生成回答失败: {str(e)}")
                    print(f"❌ 生成回答失败: {str(e)}")

            # 批量添加模式
            elif choice == '2':
                print("\n📁 输入格式说明：")
                print("- 多个路径用分号分隔")
                print("- 支持文件和文件夹路径")
                print("- 示例: D:/文档; E:/数据文件夹")

                input_paths = input("\n请输入路径：").split(';')
                paths = [p.strip() for p in input_paths if p.strip()]

                if not paths:
                    print("⚠️ 未输入有效路径")
                    continue

                print("\n🚀 开始处理文档...")
                results = batch_add_documents(paths)

                # 显示统计结果
                print("\n" + "=" * 40)
                print(f"📊 完成处理 {results['total_files']} 个文件")
                print(f"✅ 成功: {results['success_files']}")
                print(f"❌ 失败: {results['failed_files']}")
                print(f"📚 新增知识块: {results['total_chunks']}")

                # 显示错误详情
                if results['errors']:
                    print("\n详细错误报告：")
                    for error in results['errors']:
                        print(f"\n文件：{error['file']}")
                        for err in error.get('errors', []):
                            print(f"• 页{err.get('page', 'unknown')} - {err.get('error', '未知错误')}")
                            if 'sample' in err:
                                print(f"  内容示例: {err['sample']}")

            # 查看知识库状态
            elif choice == '3':
                info = vectorstore.client.get_collections()
                print("\n知识库状态：")
                print(f"• 集合数量: {len(info.collections)}")
                print(f"• 当前集合: {info.collections[0].name}")
                print(f"• 数据量: {info.collections[0].points_count} 条")

            # 退出系统
            elif choice == '4':
                print("\n感谢使用！")
                break

            else:
                print("⚠️ 无效输入，请重试")

        except KeyboardInterrupt:
            print("\n操作已取消")
            break
        except Exception as e:
            logger.exception(f"未预期错误: {str(e)}")
            print(f"❌ 发生未预期错误: {str(e)}")


if __name__ == "__main__":
    print("🔍 正在初始化知识库...")
    interactive_interface()